from pyflink.common import SimpleStringSchema, WatermarkStrategy, Types
from pyflink.datastream import StreamExecutionEnvironment, CheckpointingMode, ExternalizedCheckpointCleanup, \
    HashMapStateBackend, EmbeddedRocksDBStateBackend, KeyedProcessFunction, RuntimeContext
from pyflink.datastream.connectors.kafka import KafkaSource, KafkaOffsetsInitializer
from pyflink.datastream.state import ValueStateDescriptor, ListStateDescriptor, ListState

# 1、创建flink执行环境
env = StreamExecutionEnvironment.get_execution_environment()

env.set_parallelism(1)

source = KafkaSource.builder() \
    .set_bootstrap_servers("master:9092") \
    .set_topics("students") \
    .set_group_id("my-group") \
    .set_starting_offsets(KafkaOffsetsInitializer.latest()) \
    .set_value_only_deserializer(SimpleStringSchema()) \
    .build()

liens_ds = env.from_source(source, WatermarkStrategy.no_watermarks(), "Kafka Source")

# 实时统计每隔班级的平均年龄
kv_ds = liens_ds.map(lambda line: (line.split(",")[-1], int(line.split(",")[2])))

key_by_ds = kv_ds.key_by(lambda kv: kv[0])

class AvgAgeKeyedProcessFunction(KeyedProcessFunction):

    def __init__(self):
        self.ages_state = None

    def open(self, runtime_context: RuntimeContext):
        self.ages_state = runtime_context.get_list_state(ListStateDescriptor('ages', Types.INT()))

    def process_element(self, value, ctx: 'KeyedProcessFunction.Context'):
        clazz = value[0]
        age = value[1]

        self.ages_state.add(age)

        sum_age = 0
        num = 0
        for age in self.ages_state.get():
            sum_age += age
            num += 1

        avg_age = sum_age / num

        yield clazz, avg_age
key_by_ds.process(AvgAgeKeyedProcessFunction())


env.execute()